from typing import Dict, List
from fastapi import WebSocket
from datetime import datetime

class ConnectionManager:
    def __init__(self):
        # 存储所有活动的websocket连接
        self.active_connections: Dict[str, WebSocket] = {}
        # 存储连接时间
        self.connection_times: Dict[str, datetime] = {}
    
    async def connect(self, websocket: WebSocket, client_id: str):
        """建立连接"""
        await websocket.accept()
        self.active_connections[client_id] = websocket
        self.connection_times[client_id] = datetime.now()
    
    async def disconnect(self, client_id: str):
        """断开连接"""
        if client_id in self.active_connections:
            del self.active_connections[client_id]
            del self.connection_times[client_id]
    
    async def send_personal_message(self, message: str, client_id: str):
        """发送个人消息"""
        if client_id in self.active_connections:
            await self.active_connections[client_id].send_text(message)
    
    async def broadcast(self, message: str, exclude: str = None):
        """广播消息"""
        for client_id, connection in self.active_connections.items():
            if client_id != exclude:
                await connection.send_text(message)
    
    def get_connected_clients(self) -> List[str]:
        """获取所有已连接的客户端ID"""
        return list(self.active_connections.keys()) 